package com.luych.toolbox.common.service.tool;

import com.jcraft.jsch.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;
import org.springframework.util.StreamUtils;

import java.io.*;
import java.util.Collections;
import java.util.List;
import java.util.function.Consumer;

public class Cmd {

    private static final Logger LOGGER = LoggerFactory.getLogger(Cmd.class);

    public static String executeLocal(String cmd)  throws IOException, InterruptedException {
        return executeLocal(cmd, LOGGER::info, LOGGER::warn);
    }

    public static String executeLocal(String cmd, Consumer<String> stdout, Consumer<String> stderr) throws IOException, InterruptedException {
        return executeLocal(cmd, stdout, stderr, Collections.singletonList(0));
    }

    public static String executeLocal(String cmd, Consumer<String> stdout, Consumer<String> stderr, List<Integer> codes) throws IOException, InterruptedException {
        LOGGER.info("command for local: {}", cmd);
        cmd = "source /etc/profile; " + cmd;
        Process process = Runtime.getRuntime().exec(cmd);
        StringBuilder builder = new StringBuilder();
        Thread thread1 = logInputStream(process.getInputStream(), s -> { builder.append(s).append("\n"); stdout.accept(s);});
        Thread thread2 = logInputStream(process.getErrorStream(), s -> { builder.append(s).append("\n"); stderr.accept(s); });
        int code = process.waitFor();
        thread1.interrupt();
        thread2.interrupt();
        Assert.isTrue(codes.contains(code), "command exec error!!! return code is " + code + "\n\n" + builder);
        LOGGER.info("command exec for local success!!!");
        return builder.toString();
    }

    public static String executeRemote(String host, int port, String user, String pass, String cmd) throws JSchException, IOException, InterruptedException {
        return executeRemote(host, port,user, pass, cmd, LOGGER::info, LOGGER::warn);
    }

    public static String executeRemote(String host, int port, String user, String pass, String cmd, Consumer<String> stdout, Consumer<String> stderr) throws JSchException, IOException, InterruptedException {
        return executeRemote(host, port, user, pass, cmd, stdout, stderr, Collections.singletonList(0));
    }

    public static String executeRemote(String host, int port, String user, String pass, String cmd, List<Integer> codes) throws JSchException, IOException, InterruptedException {
        return executeRemote(host, port, user, pass, cmd, LOGGER::info, LOGGER::warn, codes);
    }

    public static String executeRemote(String host, int port, String user, String pass, String cmd, Consumer<String> stdout, Consumer<String> stderr, List<Integer> codes) throws JSchException, IOException, InterruptedException {
        Session session = null;
        ChannelExec channel = null;
        try {
            LOGGER.info("command for {}: {}", host, cmd);
            cmd = "source /etc/profile; " + cmd;
            session = new JSch().getSession(user, host, port);
            session.setConfig("StrictHostKeyChecking", "no");
            session.setPassword(pass);
            session.connect();
            channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(cmd);
            channel.connect();
            StringBuilder builder = new StringBuilder();
            Thread thread1 = logInputStream(channel.getInputStream(), s -> { builder.append(s).append("\n"); stdout.accept(s); });
            Thread thread2 = logInputStream(channel.getExtInputStream(), s -> { builder.append(s).append("\n"); stderr.accept(s); });
            while (!channel.isEOF()) {
                Thread.sleep(1000);
            }
            thread1.interrupt();
            thread2.interrupt();
            int code = channel.getExitStatus();
            Assert.isTrue(codes.contains(code), "command exec error!!! return code is " + code + "\n\n" + builder);
            LOGGER.info("command exec for {} success!!!", host);
            return builder.toString();
        } finally {
            if (channel != null && channel.isConnected()) channel.disconnect();
            if (session != null && session.isConnected()) session.disconnect();
        }
    }

    public static void uploadRemote(String host, int port, String user, String pass, InputStream inputStream, String file) throws JSchException, SftpException {
        uploadRemote(host, port, user, pass, inputStream, file, new DefaultSftpProgressMonitor());
    }

    public static void uploadRemote(String host, int port, String user, String pass, InputStream inputStream, String file, SftpProgressMonitor monitor) throws JSchException, SftpException {
        Session session = null;
        ChannelSftp channel = null;
        try {
            session = new JSch().getSession(user, host, port);
            session.setConfig("StrictHostKeyChecking", "no");
            session.setPassword(pass);
            session.connect();
            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();
            channel.put(inputStream, file, monitor);
        } finally {
            if (channel != null && channel.isConnected()) channel.disconnect();
            if (session != null && session.isConnected()) session.disconnect();
        }
    }

    public static void downloadRemote(String host, int port, String user, String pass, String file, OutputStream outputStream) throws JSchException, SftpException, IOException {
        downloadRemote(host, port, user, pass, file, outputStream, new DefaultSftpProgressMonitor());
    }

    public static void downloadRemote(String host, int port, String user, String pass, String file, OutputStream outputStream, SftpProgressMonitor monitor) throws JSchException, SftpException, IOException {
        Session session = null;
        ChannelSftp channel = null;
        try {
            session = new JSch().getSession(user, host, port);
            session.setConfig("StrictHostKeyChecking", "no");
            session.setPassword(pass);
            session.connect();
            channel = (ChannelSftp) session.openChannel("sftp");
            channel.connect();
            InputStream inputStream = channel.get(file, monitor);
            StreamUtils.copy(inputStream, outputStream);
        } finally {
            if (channel != null && channel.isConnected()) channel.disconnect();
            if (session != null && session.isConnected()) session.disconnect();
        }
    }

    private static String readLine(BufferedReader reader) {
        try {
            return reader.readLine();
        } catch (IOException e) {
            return "read error: " + e.getMessage();
        }
    }

    private static Thread logInputStream(InputStream stream, Consumer<String> log) throws IOException {
        Thread thread = new Thread(() -> {
            BufferedReader reader = new BufferedReader(new InputStreamReader(stream));
            String line = readLine(reader);
            while (line != null) {
                log.accept(line);
                line = readLine(reader);
                if (Thread.interrupted()) break;
            }
        });
        thread.start();
        return thread;
    }

    private static class DefaultSftpProgressMonitor implements SftpProgressMonitor {

        private long max;
        private String src;
        private String dest;

        @Override
        public void init(int op, String src, String dest, long max) {
            this.src = src;
            this.dest = dest;
            this.max = max;
            LOGGER.info("start transfer file from {} to {}, size {}", src, dest, max);
        }

        @Override
        public boolean count(long count) {
            float per = Math.round(100F * count / max) / 100F;
            LOGGER.info("file {}, {}/{}, {}%", src, count, max, per);
            return false;
        }

        @Override
        public void end() {
            LOGGER.info("end transfer file from {} to {}, size {}", src, dest, max);
        }
    }
}
